package com.cyf.common.core.io.fileLimit;

import java.util.concurrent.TimeUnit;

/**
 * 限流器
 * @author chenyifan
 * @create 2024-04-28 17:29
 */
public class FileSpeedLimiter {

    /**
     * KB 代表的字节数
     */
    private static final Long KB = 1024L;

    /**
     * 数据块大小，设置为 1MB
     */
    private static final Long CHUNK_LENGTH = 1024 * KB;

    /**
     * 已经发送或读取的字节数
     */
    private int haveTransferredBytes = 0;

    /**
     * 完成上一个数据块传输的时刻--单位纳秒
     */
    private long lastTransferFinishTime = System.nanoTime();

    /**
     * 允许的最大速率，默认为 1024 KB/s
     */
    private int maxRate = 1024;

    /**
     * 在最大速率下，传输一个数据块所需时间--单位纳秒
     */
    private long timeCostInMaxRate = CHUNK_LENGTH / (this.maxRate * KB) * 1000_000_000L;

    public FileSpeedLimiter() {
    }

    public FileSpeedLimiter(int maxRate) {
        if (maxRate < 0) {
            throw new RuntimeException("最大速率不能小于零");
        }
        this.maxRate = maxRate;
        if (maxRate == 0) {
            this.timeCostInMaxRate = 0;
        } else {
            this.timeCostInMaxRate = CHUNK_LENGTH / (this.maxRate * KB) * 1000_000_000L;
        }
    }

    public synchronized void limitNextBytes() {
        this.limitNextBytes(1);
    }


    /**
     * 当已经读取或写入的字节数量达到 CHUNK_LENGTH（数据块） 长度时，根据本次传输消耗时间
     * 判断是否需要睡眠，用于保证限流
     * @param len 读取或写入的字节长度
     */
    public synchronized void limitNextBytes(int len) {
        this.haveTransferredBytes += len;
        while (this.haveTransferredBytes > CHUNK_LENGTH) {
            long now = System.nanoTime();
            // 本次传输一个数据块花费的时间
            long passTime = now - this.lastTransferFinishTime;
            // 计算需要睡眠的时间
            long missedTime = this.timeCostInMaxRate - passTime;
            if (missedTime > 0) {
                try {
                    TimeUnit.NANOSECONDS.sleep(missedTime);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
            this.haveTransferredBytes -= CHUNK_LENGTH;
            this.lastTransferFinishTime = System.nanoTime();
        }
    }

}
